package com.fanshuai.algorithms.divide;

/**
 * 求第k小的数
 */
public class LastKMin {
    /**
     * 使用堆排序求解  构造极小堆
     * @param a
     * @param k
     * @return
     */
    public static int lastKMin(int[] a, int k) {
        int[] a0 = new int[a.length];
        System.arraycopy(a, 0, a0, 0, a.length);

        int n = (a0.length - 2) / 2;
        //构造极小堆
        for (int i = n; i >= 0; i--) {
            sink(a0, i, a0.length - 1);
        }

        n = a0.length - 1;
        for (int i = 1; i <= k && i <= a0.length - 1; i++) {
            swap(a0, 0, n);
            n--;
            sink(a0, 0, n);
        }

        if (k > a.length) {
            return a0[0];
        }
        return a0[a0.length - k];
    }

    private static void print(int[] a) {
        for (int i = 0; i < a.length; i++) {
            System.out.print(a[i] + " ");
        }
        System.out.print("\n");
    }

    private static void sink(int[] a, int i, int n) {
        while (2 * i + 1 <= n) {
            int k = 2 * i + 1;
            if (2 * i + 1 < n && a[k + 1] < a[k]) {
                k = k + 1;
            }

            if (a[k] < a[i]) {
                swap(a, i, k);
                i = k;
            } else {
                break;
            }
        }
    }

    private static void swap(int[] a, int i, int j) {
        int t = a[i];
        a[i] = a[j];
        a[j] = t;
    }

    /**
     * 第k小的数，使用快排变体解法：
     * partition函数将数组a分为两个子数组 left数组元素小于a[partition], right数组元素大于a[partition];
     * 如果left数组元素数量等于k-1, 返回a[partition]；
     * 如果left数组元素数量大于k-1, 在left数组中寻找第k小的元素
     * 如果left数组元素数量小于k-1, 在right数组中寻找 k - leftNum - 1小的元素
     * @param a
     * @param k
     * @return
     */
    private static int lastKMin2(int[] a, int k) {
        int[] a0 = new int[a.length];
        System.arraycopy(a, 0, a0, 0, a.length);

        return lastKMin(a0, k, 0, a0.length - 1);
    }

    private static int lastKMin(int[] a, int k, int left, int right) {
        if (left >= right) {
            return a[left];
        }

        int p = partition(a, left, right);
        int leftNum = p - left;
        if (leftNum == k - 1) {
            return a[p];
        }

        if (leftNum > k - 1) {
            return lastKMin(a, k, left, p - 1);
        } else {
            return lastKMin(a, k - leftNum - 1, p + 1, right);
        }
    }

    private static int partition(int[] a, int left, int right) {
        int base = a[right];
        int i = left - 1, j = right;

        while (i < j) {
            while (i < right && a[++i] < base) {
                ;
            }
            while (j > left && a[--j] > base) {
                ;
            }
            if (i >= j) {
                break;
            }
            swap(a, i, j);
        }
        swap(a, i, right);

        return i;
    }

    private static int partition2(int[] a, int left, int right) {
        int base = a[left];
        int i = left, j = right + 1;

        while (true) {
            while (i < right && a[++i] < base) {
                ;
            }
            while (j > left && a[--j] > base) {
                ;
            }
            if (i >= j) {
                break;
            }
            swap(a, i, j);
        }
        swap(a, j, left);

        return j;
    }

    public static void main(String[] args) {
        int[] a = {-2, -1, -6, 11, -4, 13, -5, -3};

        System.out.println(lastKMin(a, 1));
        System.out.println(lastKMin(a, 2));
        System.out.println(lastKMin(a, 3));
        System.out.println(lastKMin(a, 5));
        System.out.println(lastKMin(a, 7));
        System.out.println(lastKMin(a, 9));

        System.out.println();
        System.out.println(lastKMin2(a, 1));
        System.out.println(lastKMin2(a, 2));
        System.out.println(lastKMin2(a, 3));
        System.out.println(lastKMin2(a, 5));
        System.out.println(lastKMin2(a, 7));
        System.out.println(lastKMin2(a, 9));

    }

    private static void test() {
        int[] a = {-2, -1, -6, 11, -4, 13, -5, -3};
        int p = partition2(a, 0, a.length - 1);
        print(a);

        partition2(a, p + 1, 7);
        print(a);
    }
}
